const { PromptTemplate } = require('@langchain/core/prompts');
const { SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate } = require('@langchain/core/prompts');

// 定义系统消息模板
const SYSTEM_MESSAGES = {
  chat: "我是一个专业的AI助手，我会提供准确、有帮助的回答，并保持友好和专业的态度。",
  qa: "我是一个专业的问答助手，我会基于提供的上下文准确回答问题，如果问题超出上下文范围，我会明确指出。",
  summarize: "我是一个专业的文本摘要助手，我会准确提取文本中的关键信息，生成简洁而全面的摘要。"
};

// 定义常用的提示词模板
const TEMPLATES = {
  // 通用聊天模板
  chat: ChatPromptTemplate.fromMessages([
    SystemMessagePromptTemplate.fromTemplate(SYSTEM_MESSAGES.chat),
    HumanMessagePromptTemplate.fromTemplate("{input}")
  ]),
  
  // 问答模板
  qa: ChatPromptTemplate.fromMessages([
    SystemMessagePromptTemplate.fromTemplate(SYSTEM_MESSAGES.qa),
    HumanMessagePromptTemplate.fromTemplate("上下文: {context}\n\n问题: {question}")
  ]),
  
  // 文本摘要模板
  summarize: ChatPromptTemplate.fromMessages([
    SystemMessagePromptTemplate.fromTemplate(SYSTEM_MESSAGES.summarize),
    HumanMessagePromptTemplate.fromTemplate("请为以下文本生成摘要：\n\n{text}")
  ])
};

class PromptTemplateService {
  // 获取格式化后的提示词
  static async formatPrompt(templateName, variables) {
    const template = TEMPLATES[templateName];
    if (!template) {
      throw new Error(`模板 ${templateName} 不存在`);
    }
    
    return await template.formatMessages(variables);
  }
  
  // 获取所有可用的模板名称
  static getAvailableTemplates() {
    return Object.keys(TEMPLATES);
  }
  
  // 添加新的模板
  static addTemplate(name, systemMessage, humanTemplate) {
    if (TEMPLATES[name]) {
      throw new Error(`模板 ${name} 已存在`);
    }
    
    TEMPLATES[name] = ChatPromptTemplate.fromMessages([
      SystemMessagePromptTemplate.fromTemplate(systemMessage),
      HumanMessagePromptTemplate.fromTemplate(humanTemplate)
    ]);
    
    return name;
  }
}

module.exports = PromptTemplateService;